import time
import uuid
from datetime import datetime

from awm_client import AWMClient

def main():
    # 初始化客户端
    client = AWMClient(
        base_url="http://localhost:8080/api/v1",
        zmq_url="tcp://localhost:5555",
        api_key="your-api-key"
    )

    try:
        # 生成工作流运行ID
        workflow_run_id = str(uuid.uuid4())
        workflow_id = "example-workflow"
        user_id = "example-user"

        print(f"开始运行工作流: {workflow_run_id}")

        # 1. 创建工作流运行记录
        workflow_run = client.create_workflow_run(workflow_run_id, workflow_id)
        print(f"创建工作流运行记录: {workflow_run}")

        # 2. 发送工作流开始事件
        client.send_workflow_started(workflow_run_id, workflow_id, user_id)
        print("发送工作流开始事件")

        # 3. 模拟节点执行
        # 3.1 节点1开始
        node1_run_id = str(uuid.uuid4())
        node1_id = "node1"
        client.send_node_started(
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id=node1_run_id,
            node_id=node1_id,
            node_type=1,
            data={"input": "node1-input"}
        )
        print(f"节点1开始执行: {node1_id}")

        # 模拟节点1执行
        time.sleep(1)

        # 3.2 节点1完成
        client.send_node_finished(
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id=node1_run_id,
            node_id=node1_id,
            node_type=1,
            data={"output": "node1-output"}
        )
        print(f"节点1执行完成: {node1_id}")

        # 3.3 节点2开始
        node2_run_id = str(uuid.uuid4())
        node2_id = "node2"
        client.send_node_started(
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id=node2_run_id,
            node_id=node2_id,
            node_type=1,
            data={"input": "node2-input"}
        )
        print(f"节点2开始执行: {node2_id}")

        # 模拟节点2执行
        time.sleep(1)

        # 3.4 节点2完成
        client.send_node_finished(
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id=node2_run_id,
            node_id=node2_id,
            node_type=1,
            data={"output": "node2-output"}
        )
        print(f"节点2执行完成: {node2_id}")

        # 4. 发送工作流完成事件
        client.send_workflow_finished(workflow_run_id, workflow_id, "succeeded")
        print("发送工作流完成事件")

        # 5. 更新工作流运行状态
        workflow_run = client.update_workflow_run(
            workflow_run_id=workflow_run_id,
            status="succeeded",
            finished_at=datetime.now()
        )
        print(f"更新工作流运行状态: {workflow_run}")

        # 6. 获取工作流运行记录
        workflow_run = client.get_workflow_run(workflow_run_id)
        print(f"获取工作流运行记录: {workflow_run}")

        # 7. 获取工作流追踪日志
        traces = client.get_workflow_traces(workflow_run_id)
        print(f"获取工作流追踪日志: {len(traces)} 条记录")
        for trace in traces:
            print(f"- {trace.event}: {trace.node_id}")

    finally:
        # 关闭客户端连接
        client.close()


if __name__ == "__main__":
    main() 